from detect import FaceMeshDetector
from detect import detect
from detect import HeadPose
import cv2
class Video_transform_BS:
    def __init__(self):
        # super()
        self.fm = FaceMeshDetector()
        self.HeadPose = HeadPose()
        self.bs = [0] * 52
        self.rm = [0,0,0]
    def Get_BS(self,frame):
        frame,bs, rm = detect(frame,self.fm)
        if bs != None:
            rm =  self.HeadPose.pose_det(rm)
        return bs,rm
    def Tranform(self,filename):
        out_file_path = filename.rsplit('.', 1)[0] + '.csv'
        out_file_path = out_file_path.replace('-', '_')
        ans = []
        if filename:
            print("选择的文件为:", filename)
            # ret , frame = self.cap.read(fileName)
            # fileName = "path/to/your/video.mp4"
            capture = cv2.VideoCapture(filename)
            # ret , frame = capture.read()
            total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
            # BSTransform = bstransform()
            print("Total frames:", total_frames)
            ret, frame = capture.read()
            # time.sleep(2)
            # print(ret)
            i = 0
            # self.layout1.addWidget(self.newButton32)
            
            while True:
                ret, frame = capture.read()
                frame = cv2.flip(frame, 1)
                if not ret:
                    break
                bs,rm = self.Get_BS(frame)

                if bs == None:
                    bs=self.bs
                    rm = self.rm
                else:
                    self.rm = rm
                    self.bs = bs
                
                ans.append(','.join(str(i) for i in (self.bs+self.rm)))
                    # self.bs = bs
                
        # 打开文件并写入列表内容
            if ans!=[]:
                with open(out_file_path, 'w') as file:
                    for item in ans:
                        file.write(str(item) + '\n')
                print("成功输出BS")
        
